package com.boge.ai.utils;

import com.boge.ai.entity.embed.EmbeddingResponse;
import com.boge.ai.entity.llm.CompletionResponse;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;
import org.apache.commons.lang3.StringEscapeUtils;
import org.apache.commons.math3.linear.RealVector;

import java.util.concurrent.TimeUnit;

/**
 * 大模型的工具类
 */
public class LLMUtils {

    private static final String BASE_URL = "https://api.openai-hk.com";
    private static final String COMPLETION_URL = BASE_URL+"/v1/chat/completions";
    private static final String EMBEDDING_URL = BASE_URL+"/v1/embeddings";
    private static final String API_KEY = "hk-w3q3id1000054953c829f33c861457133e3048fa220f1076";

    /**
     * 文本生成模型接口
     * @param prompt
     * @param model
     * @return
     */
    public static String completion(String prompt,String model){
        if(model == null || model.isEmpty()){
            model = "gpt-4o";
        }
        OkHttpClient client = new OkHttpClient().newBuilder()
                .readTimeout(20, TimeUnit.SECONDS)
                .connectTimeout(20, TimeUnit.SECONDS)
                .build();
        MediaType mediaType = MediaType.parse("application/json");

        String promptMsg = String.format("""
                {
                        "model": "%s",
                        "messages": [
                            {
                                "role": "user",
                                "content": "%s"
                            }
                        ]
                }
        """,model, StringEscapeUtils.escapeJson(prompt));

        RequestBody body = RequestBody.create(mediaType, promptMsg);
        Request request = new Request.Builder()
                .url(COMPLETION_URL)
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .addHeader("Authorization", "Bearer "+API_KEY)
                .build();
        try {
            Response response = client.newCall(request).execute();
            ObjectMapper mapper = new ObjectMapper();
            String json = response.body().string();
            // 使用 Jackson 将 JSON 字符串转换为 CompletionResponse 对象
            CompletionResponse completionResponse = mapper.readValue(json, CompletionResponse.class);
            String content = completionResponse.getChoices().get(0).getMessage().getContent();

            return content;
        }catch (Exception e){
            e.printStackTrace();
        }

        return  null;
    }

    /**
     * 实现文本转向量的公共方法
     * @param prompt
     * @param model
     * @return
     */
    public static float[] embedding(String prompt,String model){
        if(model == null || model.isEmpty()){
            model = "text-embedding-3-small";
        }
        OkHttpClient client = new OkHttpClient().newBuilder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(30, TimeUnit.SECONDS)
                .build();
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, String.format("""
                {
                  "input": "%s",
                  "model": "%s"
                }
                """,prompt,model));
        Request request = new Request.Builder()
                .url(EMBEDDING_URL)
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .addHeader("Authorization", "Bearer "+API_KEY)
                .build();
        ObjectMapper mapper = new ObjectMapper();
        try {
            Response response = client.newCall(request).execute();
            ResponseBody responseBody = response.body();
            String json = responseBody.string();
            EmbeddingResponse embeddingResponse = mapper.readValue(json, EmbeddingResponse.class);
            return embeddingResponse.getData().get(0).getEmbedding();
        }catch (Exception e){
            e.printStackTrace();
        }
        return null;
    }


    /**
     * 余弦距离 -- 越⼤越相似
     * @param a
     * @param b
     * @return
     */
    public static double cosSim(RealVector a, RealVector b) {
        return a.dotProduct(b) / (a.getNorm() * b.getNorm());
    }

    /**
     * 欧式距离 -- 越⼩越相似
     * @param a
     * @param b
     * @return
     */
    public static double l2(RealVector a, RealVector b) {
        return a.subtract(b).getNorm();
    }

    /**
     * 把float[] 转换为 double[]
     * @param floatArray
     * @return
     */
    public static double[] toDoubleArray(float[] floatArray){
        if(floatArray  == null){
            return new double[0];
        }
        double[] doubleArray = new double[floatArray.length];
        for (int i = 0; i < floatArray.length; i++) {
            doubleArray[i] = floatArray[i];
        }
        return doubleArray;
    }
}
